import torch
from .BaseRumorFramework import RumorDetection
from torch.utils.data import DataLoader

class SelfTrain(RumorDetection):
    def __init__(self, sent2vec, propagation, classifier,
                 dataselector, few_shot_set,
                 batch_size=5, grad_accum_cnt=4):
        super(SelfTrain, self).__init__(sent2vec, propagation, classifier,
                                                batch_size=batch_size, grad_accum_cnt=grad_accum_cnt)
        self.dataselector = dataselector
        self.few_shot_set = few_shot_set

    def WeakLabeling(self, data):
        data_loader = DataLoader(data, batch_size=self.batch_size, shuffle=False, collate_fn=data.collate_fn)
        preds = []
        for batch in data_loader:
            pred = self.forward(batch)
            preds.append(pred)
        weak_label = torch.cat(preds).argmax(dim=1)
        data.data_y = weak_label

    def DataSelection(self, data):
        pass

    def TransferTrain(self, old_domain, new_domain):
        pass
